#include "evdeploy/cv/detection/yolox_prepostprocessor.h"

namespace ev
{
    namespace vision
    {
        YOLOxPreprocessor::YOLOxPreprocessor()
        {
        }

        YOLOxPreprocessor::~YOLOxPreprocessor()
        {
        }

        float YOLOxPreprocessor::GetScale()
        {
            return m_scale;
        }

        EVStatus YOLOxPreprocessor::Run(cv::Mat &in_mat, cv::Mat &out_mat, int out_width, bool normalize, bool bgr2rgb)
        {
            REC_TIME(t0);
            if (out_width < 10)
            {
                EVLOG(ERROR) << "input size too small:" << out_width;
                return EV_FAIL;
            }
            m_scale = std::min(out_width / static_cast<float>(in_mat.rows), out_width / static_cast<float>(in_mat.cols));
            cv::Size new_size = cv::Size{in_mat.cols * m_scale, in_mat.rows * m_scale};
            cv::Mat resized_mat;
            cv::resize(in_mat, resized_mat, new_size);

            cv::Mat board = cv::Mat(cv::Size(out_width, out_width), CV_8UC3, cv::Scalar(114, 114, 114));

            resized_mat.copyTo(board(cv::Rect{0, 0, resized_mat.cols, resized_mat.rows}));

            if (bgr2rgb)
            {
                cv::cvtColor(board, board, cv::COLOR_BGR2RGB);
            }

            if (normalize)
            {
                board.convertTo(board, CV_32F, 1 / 255.);
            }
            else
            {
                board.convertTo(board, CV_32F);
            }

            out_mat = board.clone();
            std::vector<cv::Mat> chw_wrappers;
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data + sizeof(float) * out_width * out_width);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data + 2 * sizeof(float) * out_width * out_width);
            cv::split(board, chw_wrappers);

            REC_TIME(t1);
            EVLOG(INFO) << "YOLOxPreprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }

        void GenerateGridsStride(std::vector<int> &strides, std::vector<GridAndStride> &grid_strides, int model_resolution)
        {
            for (auto stride : strides)
            {
                int num_grid_y = model_resolution / stride;
                int num_grid_x = model_resolution / stride;
                for (int g1 = 0; g1 < num_grid_y; g1++)
                {
                    for (int g0 = 0; g0 < num_grid_x; g0++)
                    {
                        grid_strides.push_back((GridAndStride){g0, g1, stride});
                    }
                }
            }
        }

        void GenerateProposals(std::vector<GridAndStride> grid_strides, EVMatData *buffer, float prob_threshold, std::vector<BoxInfo> &objects)
        {
            const int num_anchors = grid_strides.size();
            float *feat_blob = (float *)buffer->data;
            for (int anchor_idx = 0; anchor_idx < num_anchors; anchor_idx++)
            {
                const int grid0 = grid_strides[anchor_idx].grid0;
                const int grid1 = grid_strides[anchor_idx].grid1;
                const int stride = grid_strides[anchor_idx].stride;
                const int basic_pos = anchor_idx * (buffer->dims[2]);
                float x_center = (feat_blob[basic_pos + 0] + grid0) * stride;
                float y_center = (feat_blob[basic_pos + 1] + grid1) * stride;
                float w = exp(feat_blob[basic_pos + 2]) * stride;
                float h = exp(feat_blob[basic_pos + 3]) * stride;
                float x0 = x_center - w * 0.5f;
                float y0 = y_center - h * 0.5f;

                float box_objectness = feat_blob[basic_pos + 4];
                for (int class_idx = 0; class_idx < buffer->dims[2] - 5; class_idx++)
                {
                    float box_cls_score = feat_blob[basic_pos + 5 + class_idx];
                    float box_prob = box_objectness * box_cls_score;
                    if (box_prob > prob_threshold)
                    {
                        BoxInfo obj;
                        obj.x1 = x0;
                        obj.y1 = y0;
                        obj.x2 = x0 + w;
                        obj.y2 = y0 + h;
                        obj.label = class_idx;
                        obj.score = box_prob;
                        objects.push_back(obj);
                    }
                }
            }
            NMS(objects, 0.45, false);
        }

        YOLOxPostprocessor::YOLOxPostprocessor()
        {
        }

        YOLOxPostprocessor::~YOLOxPostprocessor()
        {
        }

        EVStatus YOLOxPostprocessor::Run(EVMatData *out, std::vector<ev::vision::BoxInfo> &objects, float scale, float thresh, int model_resolution, const int img_w, 
                const int img_h, float nms_thresh, bool class_aware)
        {
            REC_TIME(t0);
            std::vector<BoxInfo> proposals;
            std::vector<int> strides = {8, 16, 32};
            std::vector<GridAndStride> grid_strides;
            GenerateGridsStride(strides, grid_strides, model_resolution);
            GenerateProposals(grid_strides, out, thresh, proposals);
            objects = proposals;
            for (int i = 0; i < proposals.size(); i++)
            {
                objects[i] = proposals[i];
                float x0 = (objects[i].x1) / scale;
                float y0 = (objects[i].y1) / scale;
                float x1 = (objects[i].x2) / scale;
                float y1 = (objects[i].y2) / scale;
                x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f);
                y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f);
                x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f);
                y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f);
                objects[i].x1 = x0;
                objects[i].y1 = y0;
                objects[i].x2 = x1;
                objects[i].y2 = y1;
            }

            NMS(objects, nms_thresh, class_aware);
            REC_TIME(t1);
            EVLOG(INFO) << "YOLOxPostprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }

    } // namespace vision
} // namespace ev